from torchline.models import META_ARCH_REGISTRY

import torch
import torch.nn as nn

__all__ = ['MyModel']

@META_ARCH_REGISTRY.register()
class MyModel(nn.Module):
    def __init__(self, cfg):
        super(MyModel, self).__init__()
        self.stem = ResBlock(2, 16, 1)

        self.down1 = nn.Sequential(
            ResBlock(16, 16, 1),
            ResBlock(16, 32, 2),
        )
        self.down2 = nn.Sequential(
            ResBlock(32, 32, 1),
            ResBlock(32, 64, 2)
        )
        self.down3 = nn.Sequential(
            ResBlock(64, 64, 1),
            ResBlock(64, 128, 2)
        )
        self.up3 = nn.Sequential(
            nn.ConvTranspose1d(128, 64, 3, 2, 1, 1),
            ResBlock(64, 64, 1),
        )
        self.up2 = nn.Sequential(
            nn.ConvTranspose1d(64, 32, 3, 2, 1, 1),
            ResBlock(32, 32, 1),
        )
        self.up1 = nn.Sequential(
            nn.ConvTranspose1d(32, 16, 3, 2, 1, 1),
            ResBlock(16, 16, 1),
        )
        self.pread_layer = ResBlock(16, 2, 1)

    def forward(self, x): # b*2*128
        out_stem = self.stem(x) # b*16*128
        out_down1 = self.down1(out_stem)   # b*32*64
        out_down2 = self.down2(out_down1) # b*64*32
        out_down3 = self.down3(out_down2) # b*128*16
        out_up3 = self.up3(out_down3) + out_down2   # b*64*32
        out_up2 = self.up2(out_up3) + out_down1    # b*32*64
        out_up1 = self.up1(out_up2)  + out_stem   # b*16*128
        out = self.pread_layer(out_up1) # b*2*128
        return out

class ResBlock(nn.Module):
    def __init__(self, n_in, n_out, stride=1):
        super(ResBlock, self).__init__()
        self.conv1 = nn.Conv1d(n_in, n_out, kernel_size=3, stride=stride, padding=1)
        self.bn1 = nn.BatchNorm1d(n_out)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv1d(n_out, n_out, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm1d(n_out)

        if stride != 1 or n_out != n_in:
            self.shortcut = nn.Sequential(
                nn.Conv1d(n_in, n_out, kernel_size=1, stride=stride),
                nn.BatchNorm1d(n_out))
        else:
            self.shortcut = None

    def forward(self, x):
        residual = x
        if self.shortcut is not None:
            residual = self.shortcut(x)
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        
        out += residual
        out = self.relu(out)
        return out

if __name__ == "__main__":
    from torchline.config import get_cfg
    cfg = get_cfg()
    cfg.merge_from_file('../config/config.yaml')
    model = MyModel(cfg)
    x = torch.rand(2,2,128)
    y = model(x)
    print(x.shape)